import utils
from lib.myhmm import MyHMM
import re
from alignment import align3, split_smart
import char_inputs 

DEBUG=False



def train2(greek_file,greeklish_file):
    """ Prepares an HMM model using the greek_file for greek words and the greeklish_file for greeklish->greek mappping.
    First, it constructs a list of aligned greeklish-greek words, using the align3 function of the alignment module.
    Then, it constructs a list of split greek words, using the smart_split method of the alignment module.
    It trains the HMM model using these two and returns it to the caller.
    """
    global DEBUG
    state_transitions = {}
    start_probabilities = {'!':0}
    observations_prob = {}
    states = ['!']
    observations = []
    print 'Reading and refining words...'
    # greek_words is a probability.FreqDist instance
    greek_words = utils.simplifyWords(utils.refineWords(utils.read_words(greek_file)))
    print 'Refined: Total %d words, %d unique' % (greek_words.N(), greek_words.B())
    n = greek_words.N()
    assert n > 0
    print 'Adding words...'
    #start with an empty HMM
    hmm = MyHMM(states, observations, start_probabilities, state_transitions,observations_prob)
    
    obs_sequences = []
    grkl_file = open(greeklish_file, 'rb')
    contents = grkl_file.read()
    if contents.startswith('\xef\xbb\xbf'):
        contents = contents[3:] #strip UTF-8 BOM       
    contents = unicode(contents, 'utf-8')
    #the greeklish file is in the format of:
    #'greeklish - greek - translation_count - '
    #the translation count isn't used.
    pattern = re.compile(r'(\w+) - (\w+) - (\d+) - ', re.UNICODE)
    print 'Aligning...',
    #print repr(contents)
    for line in contents.splitlines():
        #print repr(line)
        #print repr(line.split('-'))
        m = pattern.match(line)
        if m == None: continue
        grkl_word = m.group(1)
        
        gr_word = m.group(2) 
        gr_word = gr_word.lower().translate(char_inputs.simple_ch)
        count = int(m.group(3) )
        
        #align the two words - see comments on alignment module about this
        sequence = align3(grkl_word.lower(), gr_word)
        if sequence is None: continue
        if DEBUG: utils.printSeq2(sequence)

        obs_sequences.append((sequence, count))
    print 'Done!'
    print 'Training observations...',
    #Train, use normalization, and add new states to the HMM states
    hmm.train_observations(obs_sequences, normalize=True, addStates=True)
    print 'Done!'
    
    print 'Splitting...',
    # we use the states gathered from above to split the greek words according to the above sequences
    #NEW we add all the greek chars to avoid un-split sequences...
    state_sequences = []
    parts_list = hmm.states+[i for i in char_inputs.greek_parts if i not in hmm.states]
    parts_list.sort(key=len, reverse=True)

    if DEBUG: utils.printSeq(parts_list,True)
    for word in greek_words.samples():
        parts_gr = split_smart(word.lower(), parts_list, char_inputs.gr_split)
        if parts_gr == None: continue
        if DEBUG: utils.printSeq(parts_gr)
        # for greek words we have the word count from corpora, so we use it
        state_sequences.append((parts_gr,greek_words.count(word)))
    
    print 'Done!'
    print 'Training states...',
    hmm.train_states(state_sequences, normalize=True)
    print 'Done!'
    return hmm

